"""
    波动率因子研究
    1. 基础波动率因子
    * N个月收益率标准差 (低波因子稳定性好)
    * 特质收益率回归残差波动率因子
    * 相似收益率动量因子，相似收益率波动因子受符合条件覆盖度影响较大
    * 个股beta因子
"""
import datetime
import numpy as np
import pandas as pd
import statsmodels.api as sm
from data_resource.data_bases import engine
from utilities.utilities_func import bulk_insert_with_orm, sliding_window
from scipy.stats import spearmanr, rankdata
from statsmodels.regression.rolling import RollingOLS
from scipy.stats import kurtosis
from numpy.lib.stride_tricks import sliding_window_view
from factorAnaly.factorEngine import FactorEngine
from utilities.utilities_statics import factor_standardize, factor_neutralize


def volatility_N(group, n):
    """N个交易日窗口，每日收益率标准差"""

    def _window_volatility(i):
        if i < n - 1:
            return np.nan
        window = group.iloc[i - n + 1: i + 1]
        time_decline = window.std()
        return time_decline

    signals = np.full(len(group), np.nan, dtype=np.float32)
    results = [_window_volatility(i) for i in range(n - 1, len(group))]
    # 填充结果
    signals[n - 1:] = results
    return pd.Series(signals, index=group.index)


def fama_resid(group, _market, _fama, factor_label='volatility', window_size=60):
    """
        特质收益率波动率因子
        个股收益率 = alpha + beta*市场收益率 + s*小盘股溢价 + c*高账面价值比溢价 + epsilon
        :param group: 个股收益率, columns=['trading', 'stock_pct_chg']
        :param _market: 市场收益率, columns=['trading', 'index_pct_chg']
        :param _fama: 法码三因子, columns=['trading', 'smb', 'hml']
        :param window_size: 滚动窗口
        :param factor_label: 支持的特质收益率衍生因子，volatility-残差波动率, kurtosis-残差峰度, most-残差最大异常值, count-偏离均值1个标准差的异常值数量
    """
    if factor_label not in ['volatility', 'kurtosis', 'most', 'count']:
        raise ValueError('factor_label must be volatility or kurtosis')

    # 合并dataframe
    _df1 = pd.merge(group, _market, on='trading', how='inner')
    _df = pd.merge(_df1, _fama, on='trading', how='inner')

    if _df.empty:
        raise ValueError(f"合并后数据为空，当前股票为：{group.name}")

    if len(_df) < window_size:
        rolling_std = pd.Series(data=np.nan, index=_df.index)
    else:
        X = sm.add_constant(_df[['index_pct_chg', 'smb', 'hml']])
        y = _df['stock_pct_chg']
        model = RollingOLS(y, X, window=window_size)
        results = model.fit()

        if factor_label == 'volatility':
            rolling_std = results.mse_resid
        else:
            # 矩阵计算滚动窗口残差
            _params = results.params.dropna().values  # shape:(3641, 4)
            x_windowed = sliding_window_view(X, window_size, axis=0)  # shape:(3641, 4, 60)
            y_windowed = sliding_window_view(y, window_size)  # shape:(3641, 60)
            resids = y_windowed - np.einsum('ikj,ik->ij', x_windowed, _params) # type: ignore

            if factor_label == 'kurtosis':
                k = []
                for i in range(resids.shape[0]):
                    _r = resids[i]
                    # 计算残差峰度
                    _k = kurtosis(_r, fisher=True)  # 采用超额峰度
                    k.append(_k)
                k = [np.nan] * (window_size - 1) + k
                rolling_std = k

            elif factor_label == 'most':
                k = []
                for i in range(resids.shape[0]):
                    _r = resids[i]
                    # 计算区间内异常值最大值
                    _abnormal = abs(_r - np.median(_r))
                    k.append(_abnormal.max())
                k = [np.nan] * (window_size - 1) + k
                rolling_std = k
            else:
                # count, 超过1标准差的数量
                k = []
                for i in range(resids.shape[0]):
                    _r = resids[i]
                    # 统计超过均值+1标准差的值
                    _upper = np.mean(_r) + np.std(_r)
                    _lower = np.mean(_r) - np.std(_r)
                    _k = np.sum((_r > _upper) | (_r < _lower))
                    k.append(_k)

                k = [np.nan] * (window_size - 1) + k
                rolling_std = k

    _df['signal'] = rolling_std
    return _df


def fama(startdate: str):
    """
        法码三因子模型, 股票收益可由市场风险溢价、规模因子、价值因子来合成
        * E(Rm) - Rf: 市场风险溢价，市场组合与无风险收益率差值，使用中证全指收益率替换
        * 规模因子：反应小盘股效应，小市值股票组合收益率 - 大市值股票组合收益率
        * 高PB组合收益率：高PB股票组合收益率 - 低PB股票组合收益率
    """

    # 提取股票市值、PB数据、关联个股日涨跌幅
    _sql = f"""
    select a.trade_date, a.code, a.close, a.pb, a.total_mv, c.pct_chg / 100 as pct_chg1
    from quant_research.indicator_daily a
    inner join quant_research.market_daily_ts c on a.trade_date = c.trade_date AND a.code = c.ticker
    where a.trade_date > '{startdate}'
    order by a.trade_date;
    """

    _raw = pd.read_sql(_sql, engine)

    # 市值双重分组，PB双重分组
    _raw['market_group'] = _raw.groupby('trade_date')['total_mv'].transform(
        lambda x: pd.qcut(x, 3, labels=['small_cap', 'middle_cap', 'large_cap']))
    _raw['pb_group'] = _raw.groupby('trade_date')['pb'].transform(
        lambda x: pd.qcut(x, 3, labels=['low_pb', 'middle_pb', 'high_pb']))

    # 获取唯一日期
    unique_dates = _raw['trade_date'].unique()
    n = len(unique_dates)

    # 创建包含 NaN 的 DataFrame
    _results = pd.DataFrame({
        'trade_date': unique_dates,
        'smb': [np.nan] * n,
        'hml': [np.nan] * n,
        'smb_small': [np.nan] * n,
        'smb_large': [np.nan] * n,
        'hml_small': [np.nan] * n,
        'hml_large': [np.nan] * n
    })

    # 计算分组收益率
    _smb = _raw.groupby(['trade_date', 'market_group'], observed=True)['pct_chg1'].mean()
    smb = _smb.reset_index()

    _hml = _raw.groupby(['trade_date', 'pb_group'], observed=True)['pct_chg1'].mean()
    hml = _hml.reset_index()

    _results = _results.set_index('trade_date')
    smb = smb.set_index('trade_date')
    hml = hml.set_index('trade_date')

    _results['smb_small'] = smb[smb['market_group'] == 'small_cap']['pct_chg1']
    _results['smb_large'] = smb[smb['market_group'] == 'large_cap']['pct_chg1']
    _results['hml_small'] = hml[hml['pb_group'] == 'low_pb']['pct_chg1']
    _results['hml_large'] = hml[hml['pb_group'] == 'high_pb']['pct_chg1']

    _results['smb'] = _results['smb_small'] - _results['smb_large']
    _results['hml'] = _results['hml_large'] - _results['hml_small']

    _results = _results.reset_index()

    _results['trade_date'] = pd.to_datetime(_results['trade_date']).dt.date

    # 首次写入数据库
    # bulk_insert_with_orm(_results, 'factor_fama', engine)

    return _results


def matrix_spearnman(group, rw, hw, thread, temp):
    """
        相似因子计算
        * 计算方法：假设当前时间窗口收益率序列为rw，当前窗口以前遍历所有连续历史窗口获得n个收益率序列
            分别计算历史窗口与当前窗口的秩相关系数，取绝对值符合阈值序列；
            对符合阈值序列，取下一持有期hw累计收益率
        * 相似动量因子
    """
    xx = len(group.values)
    # 生成所有的滑动窗口
    windows = [group.iloc[i: i + 60].values for i in range(0, len(group) - rw)]

    _current = np.array(windows[1:])  # 当前窗口序列数组, 维度=(2410， 60)
    _history = np.array(windows[:-1])  # 历史窗口序列数组

    # 向量化计算spearnman相关系数
    rank_his = rankdata(_history, axis=1)
    rank_cur = rankdata(_current, axis=1)

    mean_his = np.mean(rank_his, axis=1)[:, np.newaxis]
    mean_his = mean_his.repeat(rw, axis=1)  # 按列方向扩充维度
    mean_cur = np.mean(rank_cur, axis=1)[:, np.newaxis]
    mean_cur = mean_cur.repeat(rw, axis=1)  # 按列方向扩充维度

    numerator = np.dot((rank_his - mean_his), (rank_cur - mean_cur).T)
    denom1 = np.sqrt(np.sum((rank_his - mean_his) ** 2, axis=1))
    denom2 = np.sqrt(np.sum((rank_cur - mean_cur) ** 2, axis=1))
    denominator = np.outer(denom1, denom2)

    corr_matrix = numerator / denominator  # 相关系数矩阵

    group = group.to_frame(name='pct_chg')
    group['holding_return'] = group['pct_chg'][::-1].rolling(hw).apply(lambda x: (1 + x).prod() - 1)[::-1]

    _results = []
    for current in range(len(_current)):
        _corr = corr_matrix[:, current]
        _index = np.where(np.abs(corr_matrix[:, current]) >= thread)[0]
        _index1 = _index[_index <= current]  # 获取符合条件的
        if len(_index1) == 0:
            _results.append(np.nan)
            continue

        # 获取符合条件的下一期持有收益率
        _return = []
        for i in _index1:
            _idx = i + 60 - 1
            _return.append(group['holding_return'].values[_idx])

        if temp == 'momentum':
            _results.append(np.mean(_return))
        else:
            _results.append(np.std(_return))

    results = pd.Series(np.nan, index=group.index)
    # 检查 _results 长度是否合适
    if len(_results) <= len(results):
        results.iloc[-len(_results):] = _results
    else:
        print(f"Error: _results length ({len(_results)}) is longer than available space in results ({len(results)})")
        results = pd.Series(np.nan, index=group.index[:len(_results)])
        results[:] = _results

    if len(results.values) != len(group):
        print(f"Error: _results length ({len(results)}) is longer than available space in results")
    return results


def beta(group, window=250):
    """计算滚动Beta因子（固定窗口）"""
    group = group.sort_values("trading", ascending=True)
    cov = group["stock_return"].rolling(window=window, min_periods=window).cov(group["index_return"])
    var = group["index_return"].rolling(window=window, min_periods=window).var()
    return cov / var


# CHV chainkin Volatility 波动率指标
def chv_volatility():
    pass


# RVI -- 上涨波动率/整体波动率
def rvi_volatility(data: pd.DataFrame, n: int = 10, code_col: str = 'code', optimization: str = None) -> pd.DataFrame:
    """
    RVI计算：
    - numerator = close - open
    - denominator = high - low
    - 在每个滚动窗口内按时间顺序加权求和（近的权重大）
    - 返回原始数据 + rvi 列
    """
    required_columns = [code_col, 'trading', 'close', 'open', 'high', 'low']
    missing_columns = [col for col in required_columns if col not in data.columns]
    assert not missing_columns, f"缺少必要的列: {missing_columns}"

    def replace_denominator_zero(series: pd.Series, method='large', eps=1e-6):
        """
        替代分母为0的情况
        method:
            - 'small': 用很小的数 (eps)
            - 'large': 用均值
            - 'dynamic': 用分位数下限
        """
        s = series.copy()
        if method == 'small':
            s = s.replace(0, eps)
        elif method == 'large':
            s = s.replace(0, s.mean())
        elif method == 'dynamic':
            floor = s.quantile(0.05)  # 取5%分位数
            s = s.replace(0, floor)
        return s

    def weighted_sum(x):
        # 权重按时间递增 1,2,...,len(x)
        weights = np.arange(1, len(x) + 1)
        return np.sum(x * weights)

    def _calc(group: pd.DataFrame) -> pd.DataFrame:
        df = group.copy().set_index('trading')
        df['numerator'] = df['close'] - df['open']
        df['denominator'] = df['high'] - df['low']
        df['denominator'] = replace_denominator_zero(df['denominator'])  # 除0保护，用均值波动率替代。可能会导致信号稀疏

        # 滚动窗口加权和
        num_sum = df['numerator'].rolling(window=n, min_periods=1).apply(weighted_sum, raw=True)
        den_sum = df['denominator'].rolling(window=n, min_periods=1).apply(weighted_sum, raw=True)

        df['riv'] = num_sum / den_sum
        # 计算斜率
        df['riv_slope'] = df['riv'].diff()

        if optimization == 'asymmetric':
            # 非对称分子
            df['numerator_up'] = np.where(df['numerator'] > 0, df['numerator'], 0)
            df['numerator_down'] = np.where(df['numerator'] < 0, -df['numerator'], 0)  # 记录为正值

            num_up_sum = df['numerator_up'].rolling(window=n, min_periods=1).apply(weighted_sum, raw=True)
            num_down_sum = df['numerator_down'].rolling(window=n, min_periods=1).apply(weighted_sum, raw=True)

            df['riv_up'] = num_up_sum / den_sum
            df['riv_up_slope'] = df['riv_up'].diff()
            df['riv_down'] = num_down_sum / den_sum
            df['riv_down_slope'] = df['riv_down'].diff()

        if optimization == 'vol_weight':
            df['riv_slope_vol'] = df['riv_slope'] * df['vol']
        df['riv_accel'] = df['riv_slope'].diff()
        return df.reset_index()

    def _orthogonalize(group: pd.DataFrame):
        # riv斜率与PE正交
        df = group.copy()
        if len(group) < 10:
            df['riv_slope_resid'] = np.nan
            return df
        x = sm.add_constant(df['pe'].values)
        y = df['riv_slope'].values
        model = sm.OLS(y, x).fit()
        resid = y - model.predict(x)
        df['riv_slope_resid'] = resid
        return df

    result = data.groupby(code_col, group_keys=False).apply(_calc)

    # RVI因子优化方案
    if optimization == 'pe_orthogonality':
        result = result.groupby('trading', group_keys=False).apply(_orthogonalize)
        result.reset_index(drop=True, inplace=True)
        return result[data.columns.tolist() + ['riv_slope_resid']]
    elif optimization == 'vol_weight':
        return result[data.columns.tolist() + ['riv', 'riv_slope', 'riv_accel', 'riv_slope_vol']]
    elif optimization == 'asymmetric':
        return result[data.columns.tolist() + ['riv', 'riv_slope', 'riv_accel', 'riv_up', 'riv_up_slope',
                                               'riv_down', 'riv_down_slope']]
    elif optimization is None:
        return result[data.columns.tolist() + ['riv', 'riv_slope', 'riv_accel']]

